package allocator;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.concurrent.locks.ReentrantLock;

/**
 * A simple (but quite efficient) {@link ByteBuffer} based memory allocator for
 * "small" C/C++ objects.
 * <p/>
 * The allocator accepts memory allocation requests for sizes ranging from 1
 * byte up to 4096 bytes (inclusive). The internal allocation granularity is always
 * a multiple of 4, so all returned addresses should be aligned on 4 byte boundaries.
 * This allocator is thread-safe.
 * <p/>
 * A memory address can be obtained through {@link #allocate(int)} which returns an
 * {@link AddressHandle} that embeds an address represented as a primitive Java long
 * which is available via {@link AddressHandle#getAddress()}. This address can be
 * passed to C/C++ code as a jlong and (after casting it to void*) be used as the
 * start address of a memory area of the requested size. For example, one could
 * construct a C++ object at that address using placement new.
 * <p/>
 * Once you are done with the memory area it must be deallocated through
 * {@link #deallocate(AddressHandle)}, passing the same {@link AddressHandle} which
 * had been received from {@link #allocate(int)}. The {@link AddressHandle} is
 * invalid afterwards and must not be reused.
 * <p/>
 * To be functional, this class requires the "jmem32/64" shared library to be present
 * on your java.library.path. Further, it is recommended to use the "-enableassertions"
 * switch during testing.
 * <p/>
 * The structure and algorithm of this allocator mostly rests on Alexandrescu's
 * small object allocator described in chapter 4 of "Modern C++ Design" (Andrei
 * Alexandrescu, 2001).
 * <p/>
 * Copyright (c) 2012, Stefan Zobel.
 */
public final class Allocator {

	private static final Allocator instance;

	static {
		// Note: the 32/64 bit logic may not be reliable
		boolean is64BitVm = false;
		String p = System.getProperty("sun.arch.data.model");
		if (p != null) {
			is64BitVm = p.indexOf("64") != -1;
		} else {
			p = System.getProperty("os.arch");
			if (p != null && p.indexOf("64") != -1) {
				is64BitVm = true;
			}
		}
		if (is64BitVm) {
			System.loadLibrary("jmem64");
		} else {
			System.loadLibrary("jmem32");
		}
		instance = new Allocator();
	}

	private final ReentrantLock lock = new ReentrantLock(false);
	private final HashMap<Integer, FixedAllocator> pool = new HashMap<Integer, FixedAllocator>();
	private FixedAllocator lastAllocator = null;
	private FixedAllocator lastDeallocator = null;

	private Allocator() {
	}

	public static Allocator instance() {
		return instance;
	}

	public AddressHandle allocate(final int size) {
		if (size <= 0 || size > 4096) {
			throw new IllegalArgumentException("size out of accepted range (0, 4096] : " + size);
		}

		// block size computation (round up to multiple of 4)
		final int rem = size % 4;
		final int blockSize = (rem == 0) ? size : size + (4 - rem);

		lock.lock();
		try {
			if (lastAllocator != null && lastAllocator.blockSize == blockSize) {
				return lastAllocator.allocate();
			} else {
				final Integer key = Integer.valueOf(blockSize);
				FixedAllocator allocator = pool.get(key);
				if (allocator == null) {
					allocator = new FixedAllocator(blockSize);
					pool.put(key, allocator);
				}
				lastAllocator = allocator;
				return allocator.allocate();
			}
		} finally {
			lock.unlock();
		}
	}

	public void deallocate(final AddressHandle handle) {
		if (handle == null) {
			throw new IllegalArgumentException("AddressHandle : null");
		}

		assert (handle instanceof FixedAllocator.HandleImpl);

		final Integer blockSize = ((FixedAllocator.HandleImpl) handle).blockSize;
		lock.lock();
		try {
			if (lastDeallocator != null && lastDeallocator.blockSize == blockSize.intValue()) {
				lastDeallocator.deallocate(handle);
			} else {
				final FixedAllocator deallocator = pool.get(blockSize);
				if (deallocator != null) {
					deallocator.deallocate(handle);
					lastDeallocator = deallocator;
				}
			}
		} finally {
			lock.unlock();
		}
	}

	private static final class FixedAllocator {

		private final int blockSize;
		private final int numBlocksPerChunk;

		private int lastAllocIndex = -1;
		private int lastEmptyIndex = -1;
		private Chunk lastAllocChunk = null;
		private Chunk lastEmptyChunk = null;
		private ArrayList<Chunk> chunks = new ArrayList<Chunk>();

		FixedAllocator(final int block_size) {
			if (block_size <= 0 || block_size > 4096) {
				throw new IllegalArgumentException("Block size out of range (0, 4096] : " + block_size);
			}
			if (block_size % 4 != 0) {
				throw new IllegalArgumentException("Block size not a multiple of 4 : " + block_size);
			}
			// maximum is 255 blocks
			numBlocksPerChunk = (block_size <= 16) ? 255 : (4096 / block_size);
			blockSize = block_size;
		}

		AddressHandle allocate() {
			if (lastAllocChunk == null || lastAllocChunk.blocksAvailable == 0) {
				if (lastEmptyChunk != null) {
					lastAllocChunk = lastEmptyChunk;
					lastAllocIndex = lastEmptyIndex;
					lastEmptyChunk = null;
					lastEmptyIndex = -1;
				} else {
					boolean foundExisting = false;
					int firstNullIndex = -1;
					for (int i = 0; i < chunks.size(); ++i) {
						final Chunk chunk = chunks.get(i);
						if (chunk != null) {
							if (chunk.blocksAvailable > 0) {
								lastAllocChunk = chunk;
								lastAllocIndex = i;
								foundExisting = true;
								break;
							}
						} else if (firstNullIndex == -1) {
							// remember first position without chunk for later
							firstNullIndex = i;
						}
					}
					if (!foundExisting) {
						lastAllocChunk = new Chunk(blockSize, numBlocksPerChunk);
						if (firstNullIndex == -1) {
							// no free slot, so append at back
							chunks.add(lastAllocChunk);
							lastAllocIndex = chunks.size() - 1;
						} else {
							// fill first empty slot
							chunks.set(firstNullIndex, lastAllocChunk);
							lastAllocIndex = firstNullIndex;
						}
					}
				}
			}
			assert (lastAllocChunk != null && lastAllocChunk.blocksAvailable > 0);
			return new HandleImpl(lastAllocChunk.allocate(blockSize), lastAllocIndex, blockSize);
		}

		void deallocate(final AddressHandle handle) {
			assert (handle instanceof HandleImpl);

			final int index = ((HandleImpl) handle).index;

			assert (index >= 0 && index < chunks.size());

			final Chunk chunk = chunks.get(index);
			if (chunk != null) {
				chunk.deallocate(handle.getAddress(), blockSize, numBlocksPerChunk);
				if (chunk.blocksAvailable == numBlocksPerChunk) {
					// this chunk is completely empty now
					if (lastEmptyChunk == null) {
						// we always keep around one empty chunk
						lastEmptyChunk = chunk;
						lastEmptyIndex = index;
					} else if (lastEmptyChunk != chunk) {
						// we have two empty chunks, free the one with the greater index 
						if (lastEmptyIndex > index) {
							lastEmptyChunk.disposeChunk();
							chunks.set(lastEmptyIndex, null);
							lastEmptyChunk = chunk;
							lastEmptyIndex = index;
						} else if (lastEmptyIndex < index) {
							chunk.disposeChunk();
							chunks.set(index, null);
						} else {
							throw new IllegalStateException("lastEmptyIndex == index");
						}
					} else {
						throw new IllegalStateException("lastEmptyChunk == chunk");
					}
				}
			} else {
				throw new IllegalArgumentException("index: " + index);
			}
		}

		private static final class Chunk {
			private final long base;
			private ByteBuffer mem;
			private int idxNextFree = 0;
			private int blocksAvailable;

			Chunk(final int block_size, final int numBlocks) {
				mem = ByteBuffer.allocateDirect(block_size * numBlocks);
				base = getBaseAddress(mem);
				init(numBlocks - 1, block_size);
				blocksAvailable = numBlocks; 
			}

			long allocate(final int blockSize) {
				if (blocksAvailable == 0) {
					return 0L;
				}
				final int index = idxNextFree * blockSize;
				idxNextFree = mem.get(index) & 0xFF;

				// zero out memory (for C/C++ convenience only)
				mem.put(index, (byte) 0);

				--blocksAvailable;
				return base + index;
			}

			void deallocate(final long address, final int blockSize, final int numBlocks) {
				assert (address > 0L && (address - base) <= numBlocks * blockSize - 1);

				final int index = (int) (address - base);

				assert (index % blockSize == 0);
				assert (index <= blockSize * numBlocks - blockSize);

				for (int i = index + 1; i < index + blockSize; ++i) {
					// zero out memory (for C/C++ convenience only)
					mem.put(i, (byte) 0);
				}

				// set pointer to the next free block
				mem.put(index, (byte) idxNextFree);

				idxNextFree = index / blockSize;
				++blocksAvailable;
			}

			void disposeChunk() {
				if (mem != null) {
					try {
						final sun.misc.Cleaner cleaner = ((sun.nio.ch.DirectBuffer) mem).cleaner();
						if (cleaner != null) {
							cleaner.clean();
						}
					} finally {
						mem = null;
					}
				}
			}

			private void init(final int maxIdx, final int blockSize) {
				assert (maxIdx < 255);

				// set up the "pointer to next free block" chain (a singly linked list)
				for (int i = 0; i < maxIdx; ++i) {
					mem.put(i * blockSize, (byte) (i + 1));
				}
				// last block must (!) point to an invalid index
				mem.put(maxIdx * blockSize, (byte) 255);
			}

			// get the C address of the ByteBuffer as a Java long
			private static native long getBaseAddress(final ByteBuffer mem);
		}

		private static final class HandleImpl implements AddressHandle {

			private final int index;
			private final Integer blockSize;
			private final long address;

			HandleImpl(final long address, final int index, final int blockSize) {
				assert (address > 0L && index >= 0 && blockSize % 4 == 0);
				this.index = index;
				this.address = address;
				this.blockSize = Integer.valueOf(blockSize);
			}

			@Override
			public long getAddress() {
				return address;
			}
		}
	}
}
